#include "yolov8_tensorrt.h"
#include <iostream>
#include <vector>

// ⚠️ 需要修改：COCO类别名称，根据您的模型修改
std::vector<std::string> CLASS_NAMES = {
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
    "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book",
    "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
};

int main() {

    std::string enginePath = "/home/nano/Documents/rm_f_03/yolo_tensorRT/test/src/yolov8s.engine";  // ⚠️ 需要提供：TensorRT引擎文件路径
    std::string imagePath = "/home/nano/Documents/rm_f_03/yolo_tensorRT/test/src/abc.jpg";   // ⚠️ 需要提供：输入图像路径
    
    try {
        // 初始化检测器
        YOLOv8TensorRT detector(enginePath);
        
        // 读取图像
        cv::Mat image = cv::imread(imagePath);
        if (image.empty()) {
            std::cerr << "Failed to load image: " << imagePath << std::endl;
            return -1;
        }
        
        // 执行检测
        std::vector<Detection> detections = detector.detect(image);
        
        // 绘制结果
        for (const auto& det : detections) {
            cv::rectangle(image, 
                         cv::Point(det.x1, det.y1),
                         cv::Point(det.x2, det.y2),
                         cv::Scalar(0, 255, 0), 2);
            
            std::string label = CLASS_NAMES[det.classId] + 
                               " " + std::to_string(det.confidence);
            cv::putText(image, label,
                       cv::Point(det.x1, det.y1 - 5),
                       cv::FONT_HERSHEY_SIMPLEX, 0.5,
                       cv::Scalar(0, 255, 0), 2);
        }
        
        // 显示结果
        cv::imshow("YOLOv8 Detection", image);
        cv::imwrite("output.jpg", image);  // ⚠️ 可修改：输出图像路径
        cv::waitKey(0);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return -1;
    }
    
    return 0;
}